# coding=utf-8
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, ConfusionMatrixDisplay
from sklearn.svm import SVR
from yellowbrick.cluster import KElbowVisualizer
import tushare as ts

# 设置中文字体显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
sns.set(font_scale=1.2)


# 1. 数据获取与预处理
def hist_data():
    """爬取股票数据作为示例数据来源"""
    pro = ts.pro_api('9614f273ece0d16a9b13deb3cff291ac857791bed962917297a2fcc8')  # 替换为个人token
    df = pro.daily(ts_code='600519.SH', start='20180101', end='20221111')
    df.to_csv('600519.SH.csv', index=False)
    print("数据爬取完成，已保存为600519.SH.csv")


def get_smart_home_data():
    """获取智能家居设备效率数据集（示例使用Kaggle数据集链接）"""
    # 实际使用时应从指定URL下载：https://www.kaggle.com/datasets/rabieelkharoua/predict-smart-home-device-efficiency-dataset
    # 此处模拟数据加载
    print("提示：请从以下链接下载数据集并保存为smart_home_data.csv：")
    print("https://www.kaggle.com/datasets/rabieelkharoua/predict-smart-home-device-efficiency-dataset")

    # 模拟数据结构（实际使用时替换为pd.read_csv加载真实数据）
    data = pd.DataFrame({
        '设备类型': ['恒温器', '照相机', '安防系统', '智能箱'] * 50,
        '能耗': np.random.normal(5, 2, 200),
        '响应时间': np.random.normal(3, 1, 200),
        '故障率': np.random.uniform(0, 0.1, 200),
        '用户评分': np.random.uniform(3, 5, 200)
    })
    data.to_csv('smart_home_data.csv', index=False)
    return data


def preprocess_data(data):
    """数据预处理：清洗、特征选择与转换"""
    # 处理缺失值
    data = data.dropna()
    # 选择数值型特征
    numeric_cols = data.select_dtypes(include=['float64', 'int64']).columns
    features = data[numeric_cols]
    return features, data


# 2. 降维与聚类分析
def perform_pca(features, n_components=3):
    """主成分分析降维"""
    pca = PCA(n_components=n_components)
    pca_result = pca.fit_transform(features)
    pca_df = pd.DataFrame(pca_result, columns=[f'主成分{i + 1}' for i in range(n_components)])

    # 绘制 scree plot
    plt.figure(figsize=(12, 6))
    explained_variance = pca.explained_variance_ratio_
    cumulative_variance = np.cumsum(explained_variance)
    sns.barplot(x=range(1, len(explained_variance) + 1), y=explained_variance, color='b', label='单个主成分方差贡献')
    sns.lineplot(x=range(1, len(cumulative_variance) + 1), y=cumulative_variance, color='r', marker='o',
                 label='累积方差贡献')
    plt.xlabel('主成分编号')
    plt.ylabel('方差解释率')
    plt.title('PCA方差解释率分析')
    plt.legend()
    plt.show()

    return pca_df, pca


def kmeans_clustering(pca_data, max_clusters=10):
    """K-Means聚类分析"""
    # 肘部法则确定最佳聚类数
    model = KMeans(random_state=42)
    visualizer = KElbowVisualizer(model, k=(1, max_clusters))
    visualizer.fit(pca_data)
    visualizer.show()
    optimal_k = visualizer.elbow_value_

    # 执行聚类
    kmeans = KMeans(n_clusters=optimal_k, random_state=42)
    pca_data['cluster'] = kmeans.fit_predict(pca_data)
    print(f"已完成K-Means聚类，最佳聚类数：{optimal_k}")
    return pca_data, optimal_k


def visualize_clusters(pca_data):
    """可视化聚类结果"""
    # 三维散点图
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')
    clusters = pca_data['cluster'].unique()
    colors = ['y', 'g', 'b', 'r', 'orange', 'purple', 'cyan', 'magenta'][:len(clusters)]

    for cluster, color in zip(clusters, colors):
        subset = pca_data[pca_data['cluster'] == cluster]
        ax.scatter(subset['主成分1'], subset['主成分2'], subset['主成分3'],
                   c=color, label=f'聚类{cluster}', s=50)

    ax.set_xlabel('主成分1')
    ax.set_ylabel('主成分2')
    ax.set_zlabel('主成分3')
    ax.set_title('PCA降维后聚类结果三维可视化')
    plt.legend()
    plt.show()

    # 核密度分布图（以能耗特征为例）
    if '能耗' in data.columns:
        g = sns.FacetGrid(data, col='cluster', col_wrap=3)
        g.map(sns.kdeplot, '能耗', fill=True)
        g.set_titles('聚类{col_name}')
        plt.suptitle('不同聚类的能耗分布', y=1.02)
        plt.show()


# 3. 模型训练与评估
def train_evaluation_model(features, labels):
    """训练分类模型并评估"""
    x_train, x_test, y_train, y_test = train_test_split(features, labels, test_size=0.3, random_state=42)

    # 决策树模型
    dt_model = DecisionTreeClassifier(max_depth=5, random_state=42)
    dt_model.fit(x_train, y_train)

    # 预测与评估
    y_pred = dt_model.predict(x_test)
    print("决策树模型评估结果：")
    print(f"准确率：{accuracy_score(y_test, y_pred):.4f}")
    print("\n分类报告：")
    print(classification_report(y_test, y_pred))

    # 混淆矩阵可视化
    cm = confusion_matrix(y_test, y_pred)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm)
    disp.plot(cmap=plt.cm.Blues)
    plt.title('混淆矩阵可视化')
    plt.show()

    return dt_model


def stock_prediction_demo():
    """股票预测示例（报告中附录代码）"""

    def getStockData(ts_code, bar_num, start_date, end_date):
        index_data = pd.read_csv(f'{ts_code}.csv', parse_dates=['trade_date'])
        index_data = index_data[['trade_date', 'close', 'open']].sort_values(by='trade_date')
        index_data.set_index('trade_date', inplace=True)

        # 构建特征
        for i in range(1, bar_num + 1):
            index_data[f'close_{i}'] = index_data['close'].shift(i)
        index_data = index_data[bar_num:]
        index_data = index_data[(index_data.index > start_date) & (index_data.index < end_date)]

        features = index_data[[col for col in index_data.columns if 'close' in col]]
        labels = index_data['close'].shift(-1).ffill()
        return index_data, features, labels

    def trainModel(features_train, label_train):
        model = SVR(kernel='linear')
        print("开始训练=========>")
        model.fit(features_train, label_train)
        print("<========= 训练结束")
        return model

    def predictData(model, features, labels, base_data):
        date_line = [d.strftime("%Y-%m-%d") for d in labels.index]
        next_close = list(labels)
        current_close = list(features['close'])
        next_open = list(base_data['open'].shift(-1))

        predict = model.predict(features)
        df = pd.DataFrame({
            'date': date_line, 'next_close': next_close,
            'next_predict': predict, 'close': current_close, 'next_open': next_open
        })

        # 计算策略收益
        df['position'] = np.where(df['next_predict'] > df['next_open'] * 1.002, 1, 0)
        df['PL'] = np.where(df['position'] == 1, (df['next_close'] - df['next_open']) / df['next_open'], 0)
        df['strategy'] = (df['PL'].shift(1) + 1).cumprod()
        df['baseline'] = (df['next_close'].pct_change() + 1).cumprod().fillna(1)

        # 可视化
        plt.figure(figsize=(12, 6))
        plt.plot(df['strategy'], label='策略收益')
        plt.plot(df['baseline'], label='基准收益')
        plt.legend()
        plt.title('策略与基准收益对比')
        plt.show()

    # 执行示例
    hist_data()  # 确保数据存在
    index_data, features_train, labels_train = getStockData('600519.SH', 20, '2018-01-01', '2022-12-31')
    model = trainModel(features_train, labels_train)
    index_test, features_test, labels_test = getStockData('600519.SH', 20, '2023-01-01', '2023-12-31')
    predictData(model, features_test, labels_test, index_test)


# 主函数执行流程
if __name__ == "__main__":
    # 1. 数据准备
    print("====== 阶段1：数据获取与预处理 ======")
    data = get_smart_home_data()
    features, raw_data = preprocess_data(data)
    print(f"预处理完成，特征维度：{features.shape}")

    # 2. 降维分析
    print("\n====== 阶段2：PCA降维分析 ======")
    pca_result, pca_model = perform_pca(features, n_components=3)

    # 3. 聚类分析
    print("\n====== 阶段3：聚类分析 ======")
    clustered_data, k = kmeans_clustering(pca_result)
    raw_data['cluster'] = clustered_data['cluster']  # 将聚类结果关联到原始数据
    visualize_clusters(clustered_data)

    # 4. 模型训练与评估（以聚类结果为标签进行分类）
    print("\n====== 阶段4：模型训练与评估 ======")
    if 'cluster' in raw_data.columns:
        train_evaluation_model(features, raw_data['cluster'])

    # 5. 股票预测示例（可选执行）
    print("\n====== 阶段5：股票预测示例（附录代码） ======")
    # stock_prediction_demo()  # 取消注释可运行股票预测示例